import click
from datetime import datetime
from core.classifier import TransactionClassifier
from core.analyzer import SpendingAnalyzer
from core.visualizer import FinanceVisualizer
import sqlite3
from pathlib import Path

# 确保数据目录存在并初始化数据库
Path('data').mkdir(exist_ok=True)
conn = sqlite3.connect('data/finance.db')
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS transactions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    amount REAL NOT NULL,
    merchant TEXT NOT NULL,
    category TEXT NOT NULL,
    date TEXT NOT NULL
)
''')
conn.commit()
conn.close()

@click.group()
def cli():
    """个人财务助手命令行工具"""
    pass

@cli.command()
@click.argument("amount", type=float)
@click.argument("merchant")
def add_transaction(amount, merchant):
    """添加新交易记录"""
    now = datetime.now().strftime("%Y-%m-%d %H:%M")
    classifier = TransactionClassifier()
    
    # 检查模型文件是否存在，如果不存在则使用默认分类
    try:
        category = classifier.predict(amount, now, merchant)
    except FileNotFoundError:
        click.echo("模型文件未找到，使用默认分类 '购物'")
        category = "购物"
        
    # 保存到数据库
    conn = sqlite3.connect('data/finance.db')
    cursor = conn.cursor()
    cursor.execute(
        "INSERT INTO transactions (amount, merchant, category, date) VALUES (?, ?, ?, ?)",
        (amount, merchant, category, now)
    )
    conn.commit()
    conn.close()
    
    click.echo(f"添加交易：{amount}元 @ {merchant} [{category}]")

@cli.command()
def import_transactions():
    "从CSV文件导入交易数据"
    import pandas as pd
    from pathlib import Path
    
    file_path = click.prompt("请输入CSV文件路径")
    if not Path(file_path).exists():
        click.echo(f"错误：文件 {file_path} 不存在")
        return
    
    try:
        df = pd.read_csv(file_path)
        required_columns = ['amount', 'merchant', 'category', 'date']
        if not all(col in df.columns for col in required_columns):
            click.echo(f"错误：CSV文件必须包含以下列: {', '.join(required_columns)}")
            return
    
        conn = sqlite3.connect('data/finance.db')
        df.to_sql('transactions', conn, if_exists='append', index=False)
        conn.commit()
        conn.close()
        click.echo(f"成功导入 {len(df)} 条交易记录")
    except Exception as e:
        click.echo(f"导入失败: {str(e)}")

@cli.command()
def export_transactions():
    "导出交易数据到CSV文件"
    import pandas as pd
    from pathlib import Path
    
    file_path = click.prompt("请输入导出文件路径")
    if Path(file_path).exists():
        if not click.confirm(f"文件 {file_path} 已存在，是否覆盖?"):
            return
    
    try:
        conn = sqlite3.connect('data/finance.db')
        df = pd.read_sql("SELECT * FROM transactions", conn)
        df.to_csv(file_path, index=False)
        conn.close()
        click.echo(f"成功导出 {len(df)} 条交易记录到 {file_path}")
    except Exception as e:
        click.echo(f"导出失败: {str(e)}")

@cli.command()
def generate_report():
    """生成月度报表"""
    year = datetime.now().year
    month = datetime.now().month
    analyzer = SpendingAnalyzer("data/finance.db")
    data = analyzer.get_monthly_summary(year, month)
    
    if not data:
        click.echo("本月暂无交易数据，无法生成报表。")
        return
    
    FinanceVisualizer.plot_spending_trend(data)
    click.echo("已生成月度报表：monthly_report.png")

if __name__ == "__main__":
    cli()